#!/usr/bin/env python
# _*_ coding:utf-8 _*_
from __future__ import annotations

import datetime
import decimal
from typing import Any, Dict

import pymysql
from dbutils.pooled_db import PooledDB

from autoapi.common.errors import SQLSyntaxError
from autoapi.common.log import log
from autoapi.core.get_conf import autoapi_config
from autoapi.enums.query_fetch_type import QueryFetchType
from autoapi.enums.sql_type import SqlType
from autoapi.utils.enum_control import get_enum_values


class MysqlDB:
    def __init__(self) -> None:
        self._pools: Dict[str, Any] = {}
        self._default_db = autoapi_config.MYSQL_DEFAULT
        self._enabled = autoapi_config.MYSQL_ENABLE
        
        if self._enabled:
            self._init_pools()
        else:
            log.info("数据库连接已禁用，跳过初始化")

    def _init_pools(self) -> None:
        """初始化所有数据库连接池"""
        mysql_configs = autoapi_config.MYSQL_CONFIGS

        if not mysql_configs:
            log.warning("未找到MySQL配置，将使用默认配置")
            # 保持向后兼容性
            self._pools["main_db"] = PooledDB(
                pymysql,
                host=autoapi_config.MYSQL_HOST,
                port=autoapi_config.MYSQL_PORT,
                user=autoapi_config.MYSQL_USER,
                password=autoapi_config.MYSQL_PASSWORD,
                database=autoapi_config.MYSQL_DATABASE,
                charset=autoapi_config.MYSQL_CHARSET,
                maxconnections=15,
                blocking=True,
                autocommit=False,
            )
            return

        for db_name, db_config in mysql_configs.items():
            try:
                self._pools[db_name] = PooledDB(
                    pymysql,
                    host=db_config.get("host", "localhost"),
                    port=db_config.get("port", 3306),
                    user=db_config.get("user", "root"),
                    password=db_config.get("password", ""),
                    database=db_config.get("database", ""),
                    charset=db_config.get("charset", "utf8mb4"),
                    maxconnections=15,
                    blocking=True,
                    autocommit=False,
                )
                log.info(f"✅ 数据库连接池初始化成功: {db_name}")
            except Exception as e:
                log.error(f"❌ 数据库连接池初始化失败 {db_name}: {e}")
                raise

    def get_connection(self, db_name: str | None = None):
        """获取数据库连接"""
        if db_name is None:
            db_name = self._default_db

        if db_name not in self._pools:
            raise SQLSyntaxError(
                f"数据库配置 '{db_name}' 不存在，可用配置: {list(self._pools.keys())}"
            )

        return self._pools[db_name].connection()

    def init(self, db_name: str | None = None) -> tuple[Any, pymysql.cursors.DictCursor]:
        """
        初始化连接和游标

        :param db_name: 数据库名称
        :return:
        """
        conn = self.get_connection(db_name)
        cursor = conn.cursor(cursor=pymysql.cursors.DictCursor)
        return conn, cursor

    @staticmethod
    def close(conn: Any, cursor: pymysql.cursors.DictCursor) -> None:
        """
        关闭游标和数据库连接

        :return:
        """
        cursor.close()
        conn.close()

    def query(
        self, sql: str, fetch: QueryFetchType = QueryFetchType.ALL, db_name: str | None = None
    ) -> dict | list | None:
        """
        数据库查询

        :param sql: SQL语句
        :param fetch: 查询条件; one: 查询一条数据; all: 查询所有数据
        :param db_name: 数据库名称
        :return:
        """
        conn, cursor = self.init(db_name)
        try:
            cursor.execute(sql)
            if fetch == QueryFetchType.ONE:
                query_data = cursor.fetchone()
            elif fetch == QueryFetchType.ALL:
                query_data = cursor.fetchall()
            else:
                raise SQLSyntaxError(f"查询条件 {fetch} 错误, 请使用 one / all")
        except Exception as e:
            log.error(f"执行 SQL 失败: {e}")
            raise e
        else:
            log.info(f"执行 SQL 成功: {query_data}")
            if not query_data:
                return None
            return self._format_query_result(query_data)
        finally:
            self.close(conn, cursor)

    def _format_query_result(self, query_data: Any) -> Any:
        """格式化查询结果"""
        def format_value(v):
            if isinstance(v, decimal.Decimal):
                return int(v) if v % 1 == 0 else float(v)
            elif isinstance(v, datetime.datetime):
                return str(v)
            return v

        if isinstance(query_data, dict):
            return {k: format_value(v) for k, v in query_data.items()}
        elif isinstance(query_data, (list, tuple)):
            return [{k: format_value(v) for k, v in row.items()} for row in query_data]
        return query_data

    def execute(self, sql: str, db_name: str | None = None) -> int:
        """
        执行 sql 操作

        :param sql: SQL语句
        :param db_name: 数据库名称
        :return:
        """
        conn, cursor = self.init(db_name)
        try:
            rowcount = cursor.execute(sql)
            conn.commit()
        except Exception as e:
            conn.rollback()
            log.error(f"执行 SQL 失败: {e}")
            raise e
        else:
            log.info("执行 SQL 成功")
            return rowcount
        finally:
            self.close(conn, cursor)

    def exec_case_sql(
        self,
        sql: str,
        fetch: QueryFetchType | None,
        env: str | None = None,
        db_name: str | None = None,
    ) -> dict | list | int | None:
        """
        执行用例 sql

        :param sql: SQL语句
        :param fetch: 查询类型
        :param env: 环境
        :param db_name: 数据库名称
        :return:
        """
        # 获取返回数据
        if isinstance(sql, str):
            log.info(f"执行 SQL: {sql}")
            if sql.startswith(SqlType.select):
                if fetch is None:
                    fetch = QueryFetchType.ALL
                return self.query(sql, fetch, db_name)
            else:
                return self.execute(sql, db_name)

    @staticmethod
    def sql_verify(sql: str) -> None:
        sql_types = get_enum_values(SqlType)
        if not any(sql.startswith(_) for _ in sql_types):
            raise SQLSyntaxError(f"SQL 中存在非法命令类型, 仅支持: {sql_types}")


mysql_client = MysqlDB()
